import akshare as ak
import os
import pandas as pd

def create_output_directory(stock_code):
    """
    创建输出目录。

    :param stock_code: str, 股票代码
    :return: str, 输出目录路径
    """
    dir_name = f"新浪_财务报表_{stock_code}"
    if not os.path.exists(dir_name):
        os.makedirs(dir_name)
    return dir_name

def add_stock_info_to_df(df, stock_code, report_name):
    """
    在DataFrame的每一行前添加股票代码信息。

    :param df: DataFrame, 原始数据
    :param stock_code: str, 股票代码
    :param report_name: str, 报表类型
    :return: DataFrame, 添加信息后的数据
    """
    # 创建一个包含股票代码的DataFrame
    stock_info_df = pd.DataFrame({'股票代码': [stock_code] * len(df)})
    
    # 使用pd.concat一次性合并所有列，避免DataFrame碎片化
    result_df = pd.concat([stock_info_df, df.reset_index(drop=True)], axis=1)
    
    return result_df

def fetch_and_save_financial_reports(stock_code):
    """
    获取指定股票的财务三表并保存为多级表头CSV文件。
    :param stock_code: str, 股票代码（如 "600519"）
    """
    import csv
    # 创建输出目录
    output_dir = create_output_directory(stock_code)
    # 三表及其映射表路径
    report_types = {
        "资产负债表": ("资产负债表", r"e:\\cursor\\xxakshare\\三表数据清洗\\资产负债表科目表-映射表.csv", "科目名称", "类型"),
        "利润表": ("利润表", r"e:\\cursor\\xxakshare\\三表数据清洗\\利润表科目表-映射表.csv", "科目名称", "类型"),
        "现金流量表": ("现金流量表", r"e:\\cursor\\xxakshare\\三表数据清洗\\现金流量表科目表-映射表.csv", "原始科目名称", "类型")
    }
    for name, (symbol, map_path, subject_col, group_col) in report_types.items():
        try:
            print(f"正在获取 {stock_code} 的 {name} ...")
            df = ak.stock_financial_report_sina(stock=stock_code, symbol=symbol)
            if df.empty:
                print(f"⚠️ 未获取到 {stock_code} 的 {name} 数据，请检查股票代码或网络连接。")
                continue
            # 读取映射表，获取分组和科目顺序
            df_map = pd.read_csv(map_path)
            subjects = df_map[subject_col].tolist()
            groups = df_map[group_col].tolist()
            # 只保留映射表中有的科目，顺序与映射表一致
            valid_subjects = [s for s in subjects if s in df.columns]
            valid_groups = [groups[i] for i, s in enumerate(subjects) if s in df.columns]
            # 生成多级表头
            group_row = ['股票代码'] + valid_groups
            subject_row = ['股票代码'] + valid_subjects
            # 数据行
            df_data = df[valid_subjects].copy()
            df_data.insert(0, '股票代码', stock_code)
            # 写csv
            filename = os.path.join(output_dir, f"{stock_code}_{name}.csv")
            with open(filename, 'w', encoding='utf-8-sig', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(group_row)
                writer.writerow(subject_row)
                for row in df_data.values:
                    writer.writerow(row)
            print(f"✅ 已保存 {filename}")
        except Exception as e:
            print(f"❌ 获取 {stock_code} 的 {name} 出错：{e}")


def fetch_and_save_key_indicators(stock_code):
    """
    获取指定股票的关键财务指标并保存为 CSV 文件。

    :param stock_code: str, 股票代码（如 "600004"）
    """
    # 创建输出目录
    output_dir = create_output_directory(stock_code)
    
    try:
        print(f"正在获取 {stock_code} 的关键财务指标 ...")
        df = ak.stock_financial_abstract(symbol=stock_code)

        if df.empty:
            print(f"⚠️ 未获取到 {stock_code} 的关键指标数据，请检查股票代码或网络连接。")
            return

        # 转置DataFrame
        df_t = df.T
        # 第一行为分组，第二行为具体指标名
        group_row = ['股票代码'] + list(df_t.iloc[0, 1:])
        indicator_row = ['股票代码'] + list(df_t.iloc[1, 1:])
        # 数据行
        data_rows = df_t.iloc[2:, :].reset_index(drop=True)
        data_rows.columns = indicator_row
        if '股票代码' in data_rows.columns:
            data_rows = data_rows.drop(columns=['股票代码'])
        data_rows.insert(0, '股票代码', stock_code)
        # 拼接表头和数据
        from io import StringIO
        import csv
        output = StringIO()
        writer = csv.writer(output)
        writer.writerow(group_row)
        writer.writerow(indicator_row)
        for i in range(len(data_rows)):
            writer.writerow(data_rows.iloc[i].tolist())
        # 保存到指定目录
        filename = os.path.join(output_dir, f"{stock_code}_关键指标.csv")
        with open(filename, 'w', encoding="utf-8-sig", newline='') as f:
            f.write(output.getvalue())
        print(f"✅ 已保存 {filename}")

    except Exception as e:
        print(f"❌ 获取 {stock_code} 的关键指标出错：{e}")


if __name__ == "__main__":
    # 默认股票代码
    default_code = "600330"
    
    # 用户输入股票代码，默认为600330
    user_input = input(f"请输入股票代码（默认：{default_code}）：").strip()
    stock_code = user_input if user_input else default_code

    if not stock_code.isdigit() or len(stock_code) != 6:
        print("❌ 输入无效的股票代码，请输入6位数字代码（如：600519）。")
    else:
        print(f"正在处理股票代码：{stock_code}")
        
        # 获取财务三表
        fetch_and_save_financial_reports(stock_code)

        # 获取关键指标
        fetch_and_save_key_indicators(stock_code)
        
        print(f"✅ 所有数据已保存到目录：新浪_财务报表_{stock_code}")